package org.jzyan.framework.ratelimit;

import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import lombok.SneakyThrows;
import org.redisson.api.RRateLimiter;
import org.redisson.api.RateIntervalUnit;
import org.redisson.api.RateType;
import org.redisson.api.RedissonClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.Optional;
import java.util.concurrent.TimeUnit;

/**
 * <p>
 * 限流器
 * </p>
 *
 * @author jzyan
 * @since 2023-03-23
 */
@Component
public class RateLimiterService {

    @Autowired
    private RedissonClient redissonClient;

    /**
     * ratelimiter缓存容器
     */
    private static Cache<String, Optional<RRateLimiter>> rrLimiterCache = CacheBuilder.newBuilder()
            .expireAfterWrite(5, TimeUnit.MINUTES)
            .build();

    /**
     * 通用
     *
     * @param name
     * @param rateType
     * @param rate
     * @param seconds
     * @return
     */
    public boolean tryAcquire(String name, RateType rateType, int rate, int seconds) {
        //acquire
        RRateLimiter rrLimiter = getRRateLimiter("RATELIMITER_" + name, rateType, rate, seconds);
        return rrLimiter.tryAcquire();
    }

    @SneakyThrows
    public RRateLimiter getRRateLimiter(String name, RateType rateType, int rate, int seconds) {
        //如不存在相应的limiter，进行新建
        return rrLimiterCache.get(name, () -> {
            RRateLimiter rateLimiter = redissonClient.getRateLimiter(name);
            rateLimiter.setRate(rateType, rate, seconds, RateIntervalUnit.SECONDS);
            //如果不存在，进行放置
            return Optional.of(rateLimiter);
        }).get();
    }

}
